#/bin/env bash
IMG_FILE="flash.image"
TF_IMG_FILE="tf.image"
BOOT_IMG="boot.scr"
UBOOT_FILE="u-boot-sunxi-with-spl-flash.bin"
UBOOT_TF="u-boot-sunxi-with-spl-tf.bin"
DTB_FILE="sun8i-v3s-atguigupi.dtb"
KERNEL_FILE="zImage"
ROOTFS_FILE="rootfs.squashfs"
ROOTFS_EXT="rootfs.ext2"
OVERLAYFS_FILE="data.jffs2"

# 制作镜像
function make_image() {
    set -e
    # 生成32M空文件
    dd if=/dev/zero of=$IMG_FILE bs=1M count=32
    # offset=0 写入uboot
    dd if=$UBOOT_FILE of=$IMG_FILE bs=1K seek=0 conv=notrunc
    # offset=512K 写入dtb
    dd if=$DTB_FILE of=$IMG_FILE bs=1K seek=512 conv=notrunc
    # offset=544K 写入内核
    dd if=$KERNEL_FILE of=$IMG_FILE bs=1K seek=544 conv=notrunc
    # offset=7.5M 写入rootfs
    dd if=$ROOTFS_FILE of=$IMG_FILE bs=1K seek=7680 conv=notrunc
    # offset=16M 写入data.jffs2
    dd if=$OVERLAYFS_FILE of=$IMG_FILE bs=1K seek=24576 conv=notrunc
}

# 检查flash是否存在
function ckeck_flash() {
    local RESULT="$(sudo ./sunxi-fel spiflash-info)"
    test "$RESULT" = "Manufacturer: Winbond (EFh), model: 40h, size: 33554432 bytes."
}

function flash_image() {
    make_image
    echo "镜像制作完成，开始刷写镜像"
    ckeck_flash && sudo ./sunxi-fel -p spiflash-write 0 $IMG_FILE || echo "没有找到Flash"
}

function flash_uboot() {
    echo "开始刷写uboot"
    ckeck_flash && sudo ./sunxi-fel -p spiflash-write 0 $UBOOT_FILE || echo "没有找到Flash"
}

function flash_dtb() {
    echo "开始刷写dtb"
    ckeck_flash && sudo ./sunxi-fel -p spiflash-write 524288 $DTB_FILE || echo "没有找到Flash"
}

function flash_kernel() {
    echo "开始刷写kernel"
    ckeck_flash && sudo ./sunxi-fel -p spiflash-write 557056 $KERNEL_FILE || echo "没有找到Flash"
}

function flash_rootfs() {
    echo "开始刷写rootfs"
    ckeck_flash && sudo ./sunxi-fel -p spiflash-write 7864320 $ROOTFS_FILE || echo "没有找到Flash"
}

function flash_data() {
    echo "开始刷写data"
    ckeck_flash && sudo ./sunxi-fel -p spiflash-write 25165824 $OVERLAYFS_FILE || echo "没有找到Flash"
}

function print_usage() {
    echo "flash image|uboot|dtb|kernel|rootfs|data|all"
}

function make_tf_image() {
    set -e
    local BOOT_PARTITION="boot-partition"
    # 生成16M空文件
    dd if=/dev/zero of=$BOOT_PARTITION bs=1M count=16 status=progress
    sudo mkfs.vfat $BOOT_PARTITION
    sudo mount $BOOT_PARTITION /mnt
    sudo cp $BOOT_IMG /mnt/boot.scr
    sudo cp $UBOOT_TF /mnt/u-boot-sunxi-with-spl.bin
    sudo cp $DTB_FILE /mnt/sun8i-v3s-atguigupi.dtb
    sudo cp $KERNEL_FILE /mnt/zImage
    sudo umount /mnt
    # offset=8K 写入uboot
    dd if=$UBOOT_TF of=$TF_IMG_FILE bs=1K seek=8 conv=notrunc status=progress
    # 写入boot分区
    dd if=$BOOT_PARTITION of=$TF_IMG_FILE bs=1M seek=1 conv=notrunc status=progress
    # 写入根文件系统
    dd if=$ROOTFS_EXT of=$TF_IMG_FILE bs=1M seek=17 conv=notrunc status=progress
    # 创建分区表
    echo -e "1M,16M,c\n17M,1024M,L" | sudo sfdisk $TF_IMG_FILE
    # 清理boot分区
    rm -rf $BOOT_PARTITION
}

function flash_tf() {
    make_tf_image
    echo "镜像制作完成，开始刷写镜像"
    echo "查找可移除设备"
    for dev in /sys/block/sd*; do
        if test "$(cat $dev/removable)" = "1"; then
            read -p "找到可移除设备$dev, 确定刷写?(y/n)" confirm
            if test "x$confirm" = "xy"; then
                echo "开始刷写"
                sudo dd if=$TF_IMG_FILE of=/dev/$(basename $dev) status=progress
                echo "刷写完成"
                return 0
            fi
        fi
    done
}

echo "测试sudo权限"
sudo echo "成功" || echo "无法执行sudo"

for CMD in "$@"; do
    case $CMD in
    "uboot" | "dtb" | "kernel" | "rootfs" | "data") ;;
    "image")
        flash_image
        exit 0
        ;;
    "all")
        flash_uboot
        flash_dtb
        flash_kernel
        flash_rootfs
        flash_data
        exit 0
        ;;
    "tf")
        flash_tf
        exit 0
        ;;
    *)
        print_usage
        exit 1
        ;;
    esac
done

for CMD in "$@"; do
    flash_${CMD}
done
